package com.bruin_du.searcher.indexer.core;

import com.bruin_du.searcher.indexer.aop.Timing;
import com.bruin_du.searcher.indexer.mapper.IndexDatabaseMapper;
import com.bruin_du.searcher.indexer.model.Document;
import com.bruin_du.searcher.indexer.model.InvertedRecord;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;

@Slf4j
@Component
public class IndexManager {
    private final ExecutorService executorService;
    private final IndexDatabaseMapper mapper;

    public IndexManager(ExecutorService executorService, IndexDatabaseMapper mapper) {
        this.executorService = executorService;
        this.mapper = mapper;
    }

    @Timing("构建 + 保存正排索引 —— 多线程版本")
    @SneakyThrows
    public void saveForwardIndexesConcurrent(List<Document> documentList) {
        int batchSize = 10;

        int listSize = documentList.size();
        int times = (int)Math.ceil(1.0 * listSize / batchSize);
        log.debug("一共需要 {} 批任务。", times);

        CountDownLatch latch = new CountDownLatch(times);

        for (int i = 0; i < listSize; i += batchSize) {
            int from = i;
            int to = Math.min(from + batchSize, listSize);

            Runnable task = () -> {
                List<Document> subList = documentList.subList(from, to);

                mapper.batchInsertForwardIndexes(subList);

                latch.countDown();
            };

            executorService.submit(task);
        }

        latch.await();
    }

    static class InvertedInsertTask implements Runnable {
        private final CountDownLatch latch;
        private final int batchSize;
        private final List<Document> documentList;
        private final IndexDatabaseMapper mapper;

        InvertedInsertTask(CountDownLatch latch, int batchSize, List<Document> documentList, IndexDatabaseMapper mapper) {
            this.latch = latch;
            this.batchSize = batchSize;
            this.documentList = documentList;
            this.mapper = mapper;
        }

        @Override
        public void run() {
            List<InvertedRecord> recordList = new ArrayList<>();    // 放这批要插入的数据

            for (Document document : documentList) {
                Map<String, Integer> wordToWeight = document.segWordAndCalcWeight();
                for (Map.Entry<String, Integer> entry : wordToWeight.entrySet()) {
                    String word = entry.getKey();
                    int docId = document.getDocId();
                    int weight = entry.getValue();

                    InvertedRecord record = new InvertedRecord(word, docId, weight);

                    recordList.add(record);

                    // 如果 recordList.size() == batchSize，说明够一次插入了
                    if (recordList.size() == batchSize) {
                        mapper.batchInsertInvertedIndexes(recordList);  // 批量插入
                        recordList.clear();                             // 清空 list，视为让 list.size() = 0
                    }
                }
            }

            // recordList 还剩一些，之前放进来，但还不够 batchSize 个的，所以最后再批量插入一次
            mapper.batchInsertInvertedIndexes(recordList);  // 批量插入
            recordList.clear();

            latch.countDown();
        }
    }

    @Timing("构建 + 保存倒排索引 —— 多线程版本")
    @SneakyThrows
    public void saveInvertedIndexesConcurrent(List<Document> documentList) {
        int batchSize = 10000;  // 批量插入时，最多 10000 条
        int groupSize = 50;
        int listSize = documentList.size();
        int times = (int) Math.ceil(listSize * 1.0 / groupSize);
        CountDownLatch latch = new CountDownLatch(times);

        for (int i = 0; i < listSize; i += groupSize) {
            int from = i;
            int to = Integer.min(from + groupSize, listSize);
            List<Document> subList = documentList.subList(from, to);
            Runnable task = new InvertedInsertTask(latch, batchSize, subList, mapper);
            executorService.submit(task);
        }

        latch.await();
    }
}
